package online.szlib.com.common.utils;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.UUID;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import com.nimbusds.jose.EncryptionMethod;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JOSEObjectType;
import com.nimbusds.jose.JWEAlgorithm;
import com.nimbusds.jose.JWEHeader;
import com.nimbusds.jose.JWEObject;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.JWSObject;
import com.nimbusds.jose.crypto.RSADecrypter;
import com.nimbusds.jose.crypto.RSAEncrypter;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jwt.EncryptedJWT;
import com.nimbusds.jwt.JWTClaimsSet;

@Component
public class JOSEHelper {

	private static Logger logger = LoggerFactory.getLogger(JOSEHelper.class);
	
	private RSAPublicKey publicRsaKey;
	
	private RSAPrivateKey privateRsaKey;
	
	public JOSEHelper() {
		if (publicRsaKey == null || privateRsaKey == null) {
			try {
				KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
				kpg.initialize(1024);
				KeyPair keyPair = kpg.genKeyPair();
				publicRsaKey = (RSAPublicKey)keyPair.getPublic();
				privateRsaKey= (RSAPrivateKey)keyPair.getPrivate();
			} catch (Exception ex) {
				ex.printStackTrace();
				logger.debug(" >> JOSEHelper >>> " + ex.getMessage());
			}
		}
	}
	
	public RSAEncrypter getRSAEncrypt() throws JOSEException {
		return new RSAEncrypter(publicRsaKey);
	}
	
	public RSADecrypter getRSADecrypt() throws JOSEException {
		return new RSADecrypter(privateRsaKey);
	}
	
	public JWEHeader getJWEHeader() {
		return new JWEHeader(JWEAlgorithm.RSA_OAEP, EncryptionMethod.A128GCM);
	}
	
	public JWTClaimsSet getJwtClaims() {
		JWTClaimsSet jwtClaims = new JWTClaimsSet();
		jwtClaims.setIssuer("https://book-auth-server.com");
		jwtClaims.setSubject("Mariusz");
		List<String> aud = new ArrayList<>();
		aud.add("https://my-web-app.com");
		jwtClaims.setAudience(aud);
		jwtClaims.setExpirationTime(new Date(new Date().getTime() + 1000*60*30));
		jwtClaims.setNotBeforeTime(new Date());
		jwtClaims.setIssueTime(new Date());
		jwtClaims.setJWTID(UUID.randomUUID().toString());
		return jwtClaims;
	}
	
	public void setCustomerClaim(String userID){
		getJwtClaims().setCustomClaim("userID", userID);
	}
	
	public JWEObject getJWEObject(JWEHeader header, JWTClaimsSet jwtClaims) {
		return new EncryptedJWT(header, jwtClaims);
	}

	public JWEObject getEncryptJWEObject(String userID) throws JOSEException {
		RSAEncrypter encrypter = getRSAEncrypt();
		JWTClaimsSet jwtClaims = getJwtClaims();
		jwtClaims.setCustomClaim("userID", userID);
		JWEObject jwt = new EncryptedJWT(getJWEHeader(), jwtClaims);
		jwt.encrypt(encrypter);
		return jwt;
	}
	
	public JWEObject getEncryptJWEObject() throws JOSEException {
		JWEObject jwt = new EncryptedJWT(getJWEHeader(), getJwtClaims());
		jwt.encrypt(getRSAEncrypt());
		return jwt;
	}
	
	public JWEObject getEncryptJWEObject(RSAEncrypter encrypter) throws JOSEException {
		JWEObject jwt = new EncryptedJWT(getJWEHeader(), getJwtClaims());
		jwt.encrypt(encrypter);
		return jwt;
	}
	
	public void decryptJWEObject(JWEObject jwt) throws JOSEException {
		jwt.decrypt(getRSADecrypt());
	}
	
	public JWEObject getDecryptJWEObject(JWEObject jwt, RSADecrypter decrypter) throws JOSEException {
		jwt.decrypt(decrypter);
		return jwt;
	}
	
	// JWS Option
	public RSASSASigner getSigner(RSAPrivateKey privateKey) {
		return new RSASSASigner(privateKey);
	}
	
	public RSASSAVerifier getVerifier(RSAPublicKey publicKey) {
		return new RSASSAVerifier(publicKey);
	}
	
	public JWSHeader getJWSHeader() {
		JWSHeader header = new JWSHeader(JWSAlgorithm.HS256);
	    header.setKeyID("1");
		header.setType(new JOSEObjectType("JWT"));
		return header;
	}
	
	public JWTClaimsSet getJwsClaims() {
		JWTClaimsSet claimsSet = new JWTClaimsSet();
		claimsSet.setSubject("test");
		claimsSet.setIssueTime(new Date(new Date().getTime() + 1000*60*30));
		claimsSet.setIssuer("https://test.com");
		claimsSet.setCustomClaim("scope", "authid");
		return claimsSet;
	}
	
	public void signJWSObject(JWSObject jwsObject, RSASSASigner signer) throws JOSEException {
		jwsObject.sign(signer);
	}
	
	public boolean verifyJWSObject(JWSObject jwsObject, RSASSAVerifier verifier) throws JOSEException {
		return jwsObject.verify(verifier);
	}
	
}
